package org.mx.service.server.websocket;

import java.io.Serializable;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;

import org.eclipse.jetty.websocket.api.Session;
import org.mx.TypeUtils;
import org.mx.service.server.WebsocketServerConfigBean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationContext;

import lombok.Getter;

/**
 * 描述： Websocket Session管理器<br>
 * 连接会话过滤规则可能在加载的properties配置文件中进行如下定义：<br>
 * 
 * <pre>
 *     websocket.session.filter.rules=listFilterRule,DdosFilterRule
 * </pre>
 *
 * @author John.Peng Date time 2018/3/10 上午10:07
 */
public class WsSessionManager {
	private static final Logger logger = LoggerFactory.getLogger(WsSessionManager.class);

	private final Serializable setMutex = "Set task";
	private Timer pingTimer = null, cleanTimer = null;
	private int pingCycleSec = 10, cleanCycleSec = 60;
	private boolean initialized = false;

	private ApplicationContext context;
	private WebsocketServerConfigBean websocketServerConfigBean;

	private ConcurrentHashMap<String, SessionDescriptor> sessions;
	private Set<WsSessionFilterRule> rules;
	private List<WsSessionListener> wsSessionListeners;
	private Set<String> blocks;

	/**
	 * 默认的构造函数
	 *
	 * @param context                   Sping IoC容器上下文对象
	 * @param websocketServerConfigBean WebSocket服务器配置信息对象
	 */
	public WsSessionManager(ApplicationContext context, WebsocketServerConfigBean websocketServerConfigBean) {
		super();
		this.context = context;
		this.websocketServerConfigBean = websocketServerConfigBean;
	}

	private void checkServiceChecked() {
		if (!this.initialized) {
			if (logger.isWarnEnabled()) {
				logger.warn("The WebSocket Service not be initialized.");
			}
			// TODO 下一个版本可能会增加更严格的检测，不通过可能会抛出异常，帮助应用排故
		}
	}

	/**
	 * 向会话回调列表中添加指定的监听器
	 *
	 * @param listener 会话移除监听接口
	 * @return 会话监听接口列表
	 */
	public List<WsSessionListener> addSessionListener(WsSessionListener listener) {
		checkServiceChecked();
		this.wsSessionListeners.add(listener);
		return this.wsSessionListeners;
	}

	/**
	 * 从会话回调列表中移除指定的监听器
	 *
	 * @param listener 会话监听接口
	 * @return 会话监听接口列表
	 */
	public List<WsSessionListener> removeSessionListener(WsSessionListener listener) {
		checkServiceChecked();
		this.wsSessionListeners.remove(listener);
		return this.wsSessionListeners;
	}

	/**
	 * 获取指定连接关键字的会话
	 *
	 * @param connectKey 连接关键字
	 * @return 会话
	 */
	public Session getSession(WsConnectKey connectKey) {
		checkServiceChecked();
		SessionDescriptor sessionDescriptor = sessions.get(connectKey.toString());
		return sessionDescriptor == null ? null : sessionDescriptor.getSession();
	}

	/**
	 * 获取指定连接关键字的会话描述
	 *
	 * @param connectKey 连接关键字
	 * @return 会话描述
	 */
	public SessionDescriptor getSessionDescriptor(WsConnectKey connectKey) {
		checkServiceChecked();
		return sessions.get(connectKey.toString());
	}

	/**
	 * 获取被阻止连接的IP列表
	 *
	 * @return 被阻止的IP列表
	 */
	public Set<String> getBlocks() {
		checkServiceChecked();
		return blocks;
	}

	/**
	 * 获取所有的会话
	 *
	 * @return 所有会话描述集合
	 */
	public Map<String, SessionDescriptor> getSessions() {
		checkServiceChecked();
		return sessions;
	}

	/**
	 * 判定是否有规则阻止该连接请求
	 *
	 * @param session 会话
	 * @return 返回true表示阻止连接，否则返回false
	 */
	private boolean block(Session session) {
		if (rules == null) {
			return false;
		}
		for (WsSessionFilterRule rule : rules) {
			if (rule.filter(session)) {
				// 只要有一个规则符合过滤条件，就阻止连接
				blocks.add(TypeUtils.byteArray2Ip(session.getRemoteAddress().getAddress().getAddress()));
				WsConnectKey connectKey = new WsConnectKey(session);
				if (logger.isWarnEnabled()) {
					logger.warn(String.format("The session[%s] is blocked for the %s.", connectKey.toString(),
							rule.getClass().getName()));
				}
				if (wsSessionListeners != null && !wsSessionListeners.isEmpty()) {
					wsSessionListeners.forEach(listener -> listener.blocked(connectKey.toString(), session));
				}
				return true;
			}
		}
		return false;
	}

	/**
	 * 添加一个会话
	 *
	 * @param session 会话
	 */
	public void addSession(Session session) {
		checkServiceChecked();
		final int BLOCK_ERROR_CODE = 4002;
		if (session != null) {
			synchronized (setMutex) {
				WsConnectKey connectKey = new WsConnectKey(session);
				if (block(session)) {
					session.close(BLOCK_ERROR_CODE, "Be blocked for the filter rule.");
					return;
				}
				sessions.put(connectKey.toString(), new SessionDescriptor().connect(session));
				if (logger.isDebugEnabled()) {
					logger.debug(String.format("Add a websocket session[%s] successfully.", connectKey));
				}
				if (wsSessionListeners != null && !wsSessionListeners.isEmpty()) {
					wsSessionListeners.forEach(listener -> listener.added(connectKey.toString(), session));
				}
			}
		} else {
			if (logger.isWarnEnabled()) {
				logger.warn("The session is null.");
			}
		}
	}

	/**
	 * 删除一个会话
	 *
	 * @param session 会话
	 */
	public void removeSession(Session session) {
		checkServiceChecked();
		if (session != null) {
			synchronized (setMutex) {
				WsConnectKey connectKey = new WsConnectKey(session);
				removeWsSession(connectKey);
				if (logger.isDebugEnabled()) {
					logger.debug(String.format("Remove a websocket session[%s] successfully.", connectKey));
				}
			}
		} else {
			if (logger.isWarnEnabled()) {
				logger.warn("The session is null.");
			}
		}
	}

	/**
	 * 收到对应连接关键字的pong消息
	 *
	 * @param connectKey 连接关键字
	 */
	public void pong(WsConnectKey connectKey) {
		checkServiceChecked();
		if (connectKey != null) {
			synchronized (setMutex) {
				sessions.compute(connectKey.toString(), (k, v) -> {
					if (v != null) {
						return v.pong();
					} else {
						return null;
					}
				});
			}
		}
	}

	/**
	 * 删除指定的连接
	 *
	 * @param connectStrings 连接字串集合
	 * @param code           关闭状态码
	 * @param reason         关闭原因
	 */
	public void removeWsSessions(Set<String> connectStrings, int code, String reason) {
		checkServiceChecked();
		if (connectStrings != null && !connectStrings.isEmpty()) {
			synchronized (setMutex) {
				connectStrings.forEach(connectString -> removeWsSession(connectString, code, reason));
			}
		}
	}

	/**
	 * 删除指定连接关键字的会话
	 *
	 * @param connectKey 连接关键字
	 * @see #removeWsSession(String, int, String)
	 */
	private void removeWsSession(WsConnectKey connectKey) {
		checkServiceChecked();
		removeWsSession(connectKey.toString(), 1000, "Normal close operate request.");
	}

	/**
	 * 删除指定连接关键字的会话，连接关键字形如：10.1.1.21:4582
	 *
	 * @param connectKey 连接关键字
	 * @param code       关闭状态码
	 * @param reason     关闭原因
	 */
	private void removeWsSession(String connectString, int code, String reason) {
		checkServiceChecked();
		if (sessions != null && sessions.containsKey(connectString)) {
			// 先清除缓存，防止在session.close()时被重复回调。
			SessionDescriptor sessionDescriptor = sessions.get(connectString);
			sessions.remove(connectString);
			Session session = sessionDescriptor.getSession();
			try {
				session.close(code, reason);
			} catch (Exception ex) {
				if (logger.isErrorEnabled()) {
					logger.error("Close the session fail.", ex);
				}
			}
			// 通知外部监听程序，移除了相应的连接
			if (wsSessionListeners != null && !wsSessionListeners.isEmpty()) {
				wsSessionListeners.forEach(listener -> listener.removed(connectString, session));
			}
		} else {
			if (logger.isWarnEnabled()) {
				logger.warn(String.format("The session[%s] not exist, the remove operate be ignored.", connectString));
			}
		}
	}

	/**
	 * 初始化会话管理器
	 */
	public void init() {
		this.sessions = new ConcurrentHashMap<>();
		this.rules = new HashSet<>();
		this.blocks = new HashSet<>();
		this.wsSessionListeners = new ArrayList<>();
		if (websocketServerConfigBean.isEnabled()) {
			// 配置过滤规则
			if (websocketServerConfigBean != null && websocketServerConfigBean.getWebSocketFilter() != null) {
				String[] filters = websocketServerConfigBean.getWebSocketFilter().getFilters();
				if (filters != null && filters.length > 0) {
					for (String filterStr : filters) {
						WsSessionFilterRule rule = context.getBean(filterStr, WsSessionFilterRule.class);
						rule.init(this);
						rules.add(rule);
					}
				}
			}

			// 获取可能配置的ping间隔时间
			if (websocketServerConfigBean != null) {
				pingCycleSec = websocketServerConfigBean.getPingCycleSec();
				cleanCycleSec = websocketServerConfigBean.getCleanCycleSec();
			}

			// 初始化ping定时任务
			pingTimer = new Timer();
			pingTimer.scheduleAtFixedRate(new PingTask(), 0, pingCycleSec * 1000);
			// 初始化清理定时任务
			cleanTimer = new Timer();
			cleanTimer.scheduleAtFixedRate(new CleanTask(), pingCycleSec * 3 * 1000, cleanCycleSec * 1000);
			this.initialized = true;
		}
	}

	/**
	 * 销毁会话服务器
	 */
	public void destroy() {
		this.initialized = false;
		// 清理ping定时任务
		if (pingTimer != null) {
			pingTimer.cancel();
			pingTimer.purge();
			pingTimer = null;
		}
		if (cleanTimer != null) {
			cleanTimer.cancel();
			cleanTimer.purge();
			cleanTimer = null;
		}
		// 销毁过滤规则
		rules.forEach(WsSessionFilterRule::destroy);
		rules.clear();
		rules = null;
		// 断开所有连接
		int num = 0;
		if (sessions != null) {
			synchronized (setMutex) {
				num = sessions.size();
				sessions.keySet().forEach(connectString -> {
					this.removeWsSession(connectString, 1000, "Normal close operate request from destroy().");
				});
			}
			sessions = null;
		}
		if (logger.isInfoEnabled()) {
			logger.info(String.format("Clean sessions successfully, total: %d.", num));
		}
	}

	private class CleanTask extends TimerTask {
		private final int PONG_ERROR_CODE = 4001;

		/**
		 * {@inheritDoc}
		 *
		 * @see TimerTask#run()
		 */
		@Override
		public void run() {
			if (sessions == null) {
				return;
			}
			Set<String> invalid = new HashSet<>();
			sessions.forEach((connectKey, sessionDescriptor) -> {
				if (sessionDescriptor == null) {
					return;
				}
				long pingTime = sessionDescriptor.pingTime, pongTime = sessionDescriptor.pongTime;
				if (pongTime <= 0) {
					pongTime = System.currentTimeMillis();
				}
				if (Math.abs(pongTime - pingTime) >= 3 * pingCycleSec * 1000) {
					// 超过3个ping周期没有收到pong回应，则认为该连接已经异常中断
					invalid.add(connectKey);
				}
			});
			removeWsSessions(invalid, PONG_ERROR_CODE, "No pong messages have been received for more than 3 cycles.");
			if (logger.isDebugEnabled()) {
				logger.debug(String.format("Clean invalid session successfully, total: %d.", invalid.size()));
			}
		}
	}

	/**
	 * 会话描述对象，包括：连接关键字、会话、连接时间、上次Ping时间、上次Pong时间
	 */
	@Getter
	public class SessionDescriptor {
		private WsConnectKey connectKey;
		private Session session;
		private long pingTime = System.currentTimeMillis(), pongTime = 0, connectedTime = 0;

		/**
		 * 记录连接
		 *
		 * @param session 连接会话
		 * @return 会话描述
		 */
		public SessionDescriptor connect(Session session) {
			this.connectedTime = System.currentTimeMillis();
			this.connectKey = new WsConnectKey(session);
			this.session = session;
			return this;
		}

		/**
		 * 记录ping时间
		 *
		 * @return 会话描述
		 */
		public SessionDescriptor ping() {
			this.pingTime = System.currentTimeMillis();
			return this;
		}

		/**
		 * 记录pong时间
		 *
		 * @return 会话描述
		 */
		public SessionDescriptor pong() {
			this.pongTime = System.currentTimeMillis();
			return this;
		}
	}

	/**
	 * 定时向客户端发生ping命令的任务，默认调度周期为10秒。
	 */
	private class PingTask extends TimerTask {
		/**
		 * {@inheritDoc}
		 *
		 * @see TimerTask#run()
		 */
		@Override
		public void run() {
			if (sessions == null) {
				return;
			}
			sessions.forEach((connectKey, sessionDescriptor) -> {
				try {
					if (sessionDescriptor == null) {
						return;
					}
					String key = connectKey.toString();
					sessionDescriptor.getSession().getRemote().sendPing(ByteBuffer.wrap(key.getBytes()));
					sessions.put(key, sessionDescriptor.ping());
					if (logger.isDebugEnabled()) {
						logger.debug(String.format("Send ping to session[%s] successfully.", key));
					}
				} catch (Throwable ex) {
					if (logger.isWarnEnabled()) {
						logger.warn(String.format("Send ping to session[%s] fail.", connectKey.toString()));
					}
				}
			});
		}
	}
}
